import sys, struct
import math as pm
import numpy as np
import pylab as p
import math
import csv
import matplotlib.pyplot as plt
import matplotlib as mlt
import h5py
from scipy import ndimage
import scipy.sparse.linalg
import matplotlib.patches as mpatches
from numpy import genfromtxt


TOTMEMS = 30
X = 125
Y = 125
DIM = 2

PREFIX_U = 'u_r_'
PREFIX_V = 'v_r_'
PREFIX_RHO = 'rho_r_'

#OUTPUT_DATA_DIR = '/home/behollis/Dropbox/visweek2015/oslines/'
#INPUT_DATA_DIR = '/home/behollis/DATA/in/ts00050/'
OUTPUT_DATA_DIR = '/home/brad/visweek2015-revision-code/oslines/'
INPUT_DATA_DIR = '/home/brad/DATA/ts00050/'

def readMember(mem):
    ''' Read numpy matrix from disk and return tuple of realization data. '''
    filename_u = PREFIX_U + str(mem) + '.txt'
    filename_v = PREFIX_V + str(mem) + '.txt'
    filename_rho = PREFIX_RHO + str(mem) + '.txt'
    
    u = genfromtxt(INPUT_DATA_DIR + filename_u)
    u = -1.0*u[0:X,0:Y]
    
    v = genfromtxt(INPUT_DATA_DIR + filename_v)
    v = v[0:X,0:Y]
    #rho = genfromtxt(PATH + filename_rho)
    
    return (u,v)

def showSlines2(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    vclin = np.zeros( shape = (TOTMEMS, X, Y, DIM) )
    idx = 0
    mems = [ d for d in f ]
    for mem in mems:
        uv = readMember(int(mem[3:-1])+1)
        vclin[idx][:,:,0] = uv[0]
        vclin[idx][:,:,1] = uv[1] 
        idx += 1
    
    
    mag = np.zeros(shape=(X,Y))
    for x in range(0,X):
        for y in range(0,Y):
            avg_vel_mag = 0.0
            for idx in range(7,8):
                avg_vel_mag += math.sqrt( math.pow(vclin[idx][x,y][0],2) + math.pow(vclin[idx][x,y][1],2) ) 
                
            mag[x,y] = avg_vel_mag# / float(TOTMEMS)
    
    plt.imshow(mag, origin='lower', interpolation=None, vmin=0, vmax=4, cmap='gist_ncar')
    plt.colorbar()
    
    '''
    mems = [ d for d in f ]
    for m in mems:
        x = '/x020'
        y = '/y120'
        sl = f[m+x+y]
        plt.plot(sl[0], sl[1], color ='white')
        print len(sl[0])
        
    plt.plot(20., 120., '+', color='red', linewidth=50.0, ms=10.0) # draw seed loc
    '''

    plt.show()
    
    f.close()

def showSlines(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    '''
    mag = np.zeros(shape=(X,Y))
    for x in range(0,X):
        for y in range(0,Y):
            avg_vel_mag = 0.0
            for idx in range(7,8):
                avg_vel_mag += math.sqrt( math.pow(vclin[idx][x,y][0],2) + math.pow(vclin[idx][x,y][1],2) ) 
                
            mag[x,y] = avg_vel_mag# / float(TOTMEMS)
    
    plt.imshow(mag, origin='lower', interpolation=None, vmin=0, vmax=4, cmap='gist_ncar')
    plt.colorbar()
    '''
    
    fig = plt.subplot()
    
    mems = [ d for d in f ]
    for m in mems:
        x = '/x035'
        y = '/y030'
        sl = f[m+x+y]
        plt.plot(sl[0], sl[1])
        print len(sl[0])
        
    plt.plot(35, 30, '+', color='red', linewidth=50.0, ms=10.0) # draw seed loc
    
    fig.set_xlim([0,125])
    fig.set_ylim([0,125])
    
    plt.show()
    
    f.close()

def showFTVA(file):
    
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
   
    ftva = np.zeros(shape=(X,Y))
   
    for x in range(0,X-2):
        for y in range(0, Y-2):
            c = f['cov'][x,y]
            
            #try:
            w, v = np.linalg.eig(c)
            sorted(w)
            #except:
            #    w[0] = 0
                
            ftva[x,y] = w[0]
    
    plt.imshow(ftva, origin='lower', interpolation=None, vmin=0, vmax=40)
    plt.colorbar()
    plt.show()
    f.close()

def showEntropySlineClusters(file,mult,x=None,y=None):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['slineclusters'] 
    
    max = int(np.amax(termClusters))
    
    FACT = mult
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*FACT*(max+1),max*FACT)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))

    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None, vmin=0, vmax=10)
    plt.colorbar()
    
    '''
    patches = list()    
    for idx in range(0,len(cmaplist)):
        patches.append( mpatches.Patch(color='red', label=str(idx)+' clusters') )
        
    plt.legend(handles=patches)
    '''
    
    #plot slines at seed
    if x is not None:
        f = h5py.File(OUTPUT_DATA_DIR + 'oceanSlinesLevel00.hdf5', 'r')
        mems = [ d for d in f ]
        
        for m in mems:
            dir = m + '/x' + str(y).zfill(3) + '/y' + str(x).zfill(3)
            sl = f[dir]
            plt.plot( sl[1], sl[0] )
            
        plt.plot(x, y, '+', color='white', linewidth=10.0, ms=10.0) # draw seed loc
    
    plt.show()
    f.close()

def showTerminalParticleClusters(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptclusters'] 
    
    max = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''
    
    plt.colorbar()
    
    plt.show()
    f.close()
    
def showSlineSamplingRate(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptsSampledForClustering'] 
    
    max = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''
    
    plt.colorbar()
    
    plt.show()
    f.close()
    
def showEntropyFields(file):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    
    avgAngEntr = f['avgAngularEntropy']
    avgAngEntrMin = np.min(avgAngEntr)
    avgAngEntrMax = np.max(avgAngEntr)
    avgLinEntr = f['avgLinearEntropy']
    avgLinEntrMin = np.min(avgLinEntr)
    avgLinEntrMax = np.max(avgLinEntr)
    
    
    '''
    dsampled = output.create_dataset(name='ptsSampledForClustering', shape=(LAT,LON), dtype='f')
    dentroL = output.create_dataset(name='avgAngularEntropy', shape=(LAT,LON), dtype='f')
    dentroA = output.create_dataset(name='avgLinearEntropy', shape=(LAT,LON), dtype='f')
    '''

    # angular entropy map
    cmap = plt.cm.jet
    
    plt.imshow(avgAngEntr, origin='lower', interpolation=None, vmin=0.995, vmax=1.0, cmap=cmap)
    cb = plt.colorbar(); 
    plt.show()
    
    # linear entropy map
    plt.imshow(avgLinEntr, origin='lower', interpolation=None, vmin=0.6, vmax=1.0, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    
    '''
    # sum of maps
    plt.imshow(np.add(avgLinEntr,avgAngEntr), origin='lower', interpolation=None, vmin=0.6, vmax=1.0, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    '''
    
    powOf2 = np.zeros(shape=avgLinEntr.shape); powOf2.fill(2.0)
    sqRoot = np.zeros(shape=avgLinEntr.shape); sqRoot.fill(0.5)
    
    gradLinEntr = np.gradient(avgLinEntr)
    gradLinEntrMag = np.power( np.power(gradLinEntr[0], powOf2) + np.power(gradLinEntr[1], powOf2), sqRoot)
    
    gradAngEntr = np.gradient(avgAngEntr)
    gradAngEntrMag = np.power( np.power(gradAngEntr[0], powOf2) + np.power(gradAngEntr[1], powOf2), sqRoot)
    
    # GRADIENT of LINEAR ENTROPY map
    plt.imshow(gradAngEntrMag, origin='lower', interpolation=None, vmin=np.min(gradAngEntrMag), vmax=0.05, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    
    # GRADIENT of A ENTROPY map
    plt.imshow(gradLinEntrMag, origin='lower', interpolation=None, vmin=np.min(gradLinEntrMag), vmax=0.05, cmap=cmap)
    plt.colorbar(); 
    plt.show()
    
    f.close()
    
def showSlineSamplingRate2(file, mult=10):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['ptsSampledForClustering'] 
    
    max = int(np.amax(termClusters))
    min = int(np.amax(termClusters))
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*mult*(max+1),max*mult)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))
    
    fig, ax = plt.subplots()
    
    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    bounds = bounds[3:-1]
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)#, vmin = 3, vmax = max)
    
    '''
    ps = list()
    
    for idx in range( 0,len(cmaplist) ):
        p, = mpatches.Patch(color='red', label='2' )
        ps.append(p) 
        
    plt.legend(handles=ps)
    '''
    tlabs = [str(i) for i in range(3,max+1)]
    tks = [i + 0.5 for i in range(3,max+1)]
    
    
    #plt.colorbar()
    cb = fig.colorbar(imgplot, ticks=tks)
    cb.ax.set_yticklabels(tlabs)
    
    plt.show()
    f.close()
    
def showEntropySlineClusters2(file,mult,x=None,y=None):
    f = h5py.File(OUTPUT_DATA_DIR + file, 'r')
    termClusters = f['slineclusters'] 
    
    max = int(np.amax(termClusters))
    
    FACT = mult
    
    # define the colormap
    cmap = plt.cm.jet
    # extract all colors from the .jet map
    cmaplist = [cmap(i) for i in range(0,max*FACT*(max+1),max*FACT)]
    # force the first color entry to be grey
    # create the new map
    cmap = cmap.from_list('Custom cmap', cmaplist, len(cmaplist))

    # define the bins and normalize
    bounds = np.linspace(0,max+1,max+2)
    norm = mlt.colors.BoundaryNorm(bounds, cmap.N)
    
    fig = plt.figure()
    
    imgplot = plt.imshow(termClusters, origin='lower',cmap=cmap, norm=norm, interpolation=None)
    #plt.colorbar()
    
    tlabs = [str(i) for i in range(0,max+1)]
    tks = [i + 0.5 for i in range(0,max+1)]
    
    
    #plt.colorbar()
    cb = fig.colorbar(imgplot, ticks=tks)
    cb.ax.set_yticklabels(tlabs)
    
    '''
    patches = list()    
    for idx in range(0,len(cmaplist)):
        patches.append( mpatches.Patch(color='red', label=str(idx)+' clusters') )
        
    plt.legend(handles=patches)
    '''
    
    '''
    #plot slines at seed
    if x is not None:
        f = h5py.File(OUTPUT_DATA_DIR + 'oceanSlinesLevel00.hdf5', 'r')
        mems = [ d for d in f ]
        
        for m in mems:
            dir = m + '/x' + str(y).zfill(3) + '/y' + str(x).zfill(3)
            sl = f[dir]
            plt.plot( sl[1], sl[0] )
            
        plt.plot(x, y, '+', color='white', linewidth=10.0, ms=10.0) # draw seed loc
    '''
    
    plt.show()
    f.close()
    

if __name__ == '__main__':
    
    TOTMEMS = 30
    X = 125
    Y = 125
    DIM = 2
    
#    showSlines('lockSlines.1100steps.hdf5')
#    showFTVA(file='lockFtvaForward.hdf5')
#    showFTVA(file='lockFtvaBack.hdf5') 
#    showTerminalParticleClusters(file='lockTermClusForward.hdf5', mult=15)
#    showTerminalParticleClusters(file='lockTermClusBack.hdf5', mult=15)
    showEntropySlineClusters2(file='lock.SlineClusters.varSampling.3to13pts.hdf5',mult=4)
    showSlineSamplingRate2(file='lock.SlineClusters.varSampling.3to13pts.hdf5', mult=1)
#    showEntropySlineClusters(file='lockSlineClusters.10pts.hdf5',mult=4)
#    showEntropyFields(file='lockEntropyONLY.fixed.hdf5')
    
    
    
    
    
    
    